import csv

import numpy as np
import torch

if __name__ == "__main__":
    compare_result_path = "testing.csv"
    input_dict_list = [
        {"input": np.array([[[0., 1., 2.]]], dtype=np.float32)},
        {"input": np.array([[[255., 255., 255.]]], dtype=np.float64)}
    ]

    source_results = []
    follow_results = []

    for input_dict in input_dict_list:
        # Convert Numpy to CUDA torch.Tensor
        for key in input_dict.keys():
            val = input_dict[key]
            if type(val) == np.ndarray:
                input_dict[key] = torch.from_numpy(val).cuda()
        try:
            a = torch.log(**input_dict)
            print("OP 1 output:")
            print(a)
            source_results.append(a.cpu().numpy())
        except Exception as e:
            print("[OP 1] Exception occurred~:\n", e)
            a = torch.tensor([])

        try:
            b = torch.exp(a)
            print("OP 2 output")
            print(b)
            follow_results.append(b.cpu().numpy())
        except Exception as e:
            print("[OP 2] Exception occurred~:\n", e)
            b = torch.tensor([])

    # Compare source results and follow-up results
    compare_results = []
    for i in range(len(source_results)):
        a = source_results[i]
        b = follow_results[i]
        if np.allclose(a, b, atol=1.e-2, rtol=1.e-5, equal_nan=True):
            compare_results.append([i, "Consistent"])
        else:
            compare_results.append([i, "Inconsistent", a.shape, a.dtype])
    with open(compare_result_path, "w+", newline="") as f:
        w = csv.writer(f)
        w.writerows(compare_results)



